package com.lonewalker.sparrow.util;


import com.lonewalker.sparrow.domain.TreeNode;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author lonewalker
 */
public class TreeUtil {

    public static <T extends TreeNode> List<T> buildTree(List<T> treeNodeList){
        List<T> rootNodeList = new ArrayList<>();
        //移除根数据
        rootNodeList = treeNodeList.stream().filter(d -> d.getParentId() == null || d.getParentId() == 0L).collect(Collectors.toList());
        treeNodeList.removeAll(rootNodeList);
        //根据父节点分组
        Map<Long, List<T>> childMap = treeNodeList.stream().collect(Collectors.groupingBy(T::getParentId));
        //设置子节点[这里不包括根节点]
        treeNodeList.forEach(treeNode -> treeNode.setChildren(castList(childMap.get(treeNode.getId()),TreeNode.class)));
        //将上面已经构建的子树拼到对应的根节点上
        return rootNodeList.stream().peek(root -> root.setChildren(castList(childMap.get(root.getId()),TreeNode.class))).collect(Collectors.toList());
    }

    /**
     * 这部分不是必要的，childMap.get(treeNode.getId())强转也行
     */
    public static <T> List<T> castList(Object obj, Class<T> clazz)
    {
        List<T> result = new ArrayList<T>();
        if(obj instanceof List<?>) {
            for (Object o : (List<?>) obj)
            {
                result.add(clazz.cast(o));
            }
            return result;
        }
        return null;
    }
}
